#!/usr/bin/python

# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# Implement a pseudo cdma modem.
#
# This modem mimics a CDMA modem and allows a user to experiment with
# a modem which starts in a factory reset state and is gradually moved
# into a fully activated state.
#
# To test you'll need to have a machine with at least 1 ethernet port
# available to simulate the cellular connection.  Assume that it is
# called eth1.
#   1.  install flimflam-test on your DUT.
#   2.  sudo backchannel setup eth0 pseudo-modem0
#   3.  activation-server &
#   4.  sudo fake-cromo
#   5.  Use the UI to "Activate Test Network"
#


import dbus, glib, gobject, os, subprocess, sys, time
from optparse import OptionParser

import_path = os.environ.get('SYSROOT', '/usr/local') + '/usr/lib/flimflam/test'
sys.path.append(import_path)

import flimflam_test

Modem = flimflam_test.Modem
ModemManager = flimflam_test.ModemManager

class IpPoolRestrictor:
    def __init__(self, interface):
        self.interface = interface
        self.restricted = False

    def enter(self):
        # Reject all non local tcp traffic, but allow DNS lookups
        if self.restricted:
            return
        subprocess.call(['iptables',
                         '-I', 'INPUT', '1',
                         '-p', 'tcp',
                         '-i', self.interface,
                         '-j', 'REJECT'])
        self.restricted = True

    def leave(self, force=None):
        if self.restricted or force:
            subprocess.call(['iptables',
                               '-D', 'INPUT',
                               '-p', 'tcp',
                               '-i', self.interface,
                               '-j', 'REJECT'])
            self.restricted = False

restrictor = IpPoolRestrictor('pseudo-modem0')


class CarrierState(dbus.service.Object):
    def __init__(self, bus, path):
        self.payment_made = False
        self.restricted = False
        dbus.service.Object.__init__(self, bus, path)

    @dbus.service.method('org.chromium.ModemManager.Carrier',
                         in_signature = '', out_signature = '')
    def ProcessPayment(self, *_args, **_kwargs):
        print "CarrierState: ProcessPayment"
        self.payment_made = True
        self.restricted = False

    @dbus.service.method('org.chromium.ModemManager.Carrier',
                         in_signature = '', out_signature = '')
    def ConsumePlan(self, *_args, **_kwargs):
        print "CarrierState: ConsumePlan"
        self.payment_made = False
        self.restricted = True


class FactoryResetModem(Modem):

    def __init__(self, mm, name):
        Modem.__init__(self, mm, name,
                       mdn='0000001234',
                       activation_state=Modem.NOT_ACTIVATED)

    def Activate(self, s, *_args, **_kwargs):
        print 'FactoryResetModem: Activate "%s"' % s
        self.StartActivation(Modem.PARTIALLY_ACTIVATED,
                             self.manager.MakePartiallyActivatedModem,
                             '0015551212')

    # Implement connect as a failure
    def Connect(self, _props, *_args, **_kwargs):
        print 'FactoryResetModem: Connect'
        time.sleep(self.manager.options.connect_delay_ms / 1000.0)
        self.state = flimflam_test.STATE_CONNECTING
        glib.timeout_add(500, lambda: self.ConnectDone(
                         self.state,
                         flimflam_test.STATE_REGISTERED,
                         flimflam_test.REASON_USER_REQUESTED))
        raise flimflam_test.ConnectError()

    def ActivateImpl(self, _s, _args, _kwargs):
        raise NotImplementedError('Unimplemented.  Must implement in subclass.')


class PartiallyActivatedModem(Modem):

    def __init__(self, mm, name):
        Modem.__init__(self, mm, name,
                       mdn='0015551212',
                       activation_state=Modem.PARTIALLY_ACTIVATED)

    def Activate(self, s, *_args, **_kwargs):
        print 'Partially_ActivatedModem: Activate "%s"' % s
        carrier = self.manager.carrier
        if self.manager.options.activatable and carrier.payment_made:
            self.StartActivation(Modem.ACTIVATED,
                                 self.manager.MakeActivatedModem,
                                 '6175551212')
        else:
            # TODO(jglasgow): define carrier error codes
            carrier_error = 1
            self.StartFailedActivation(carrier_error)

    def ConnectDone(self, old, new, why):
        # Implement ConnectDone by manipulating the IP pool restrictor
        if new == flimflam_test.STATE_CONNECTED:
            restrictor.enter()
        else:
            restrictor.leave()
        Modem.ConnectDone(self, old, new, why)

    def ActivateImpl(self, _s, _args, _kwargs):
        raise NotImplementedError('Unimplemented.  Must implement in subclass.')


class ActivatedModem(Modem):
    def __init__(self, mm, name):
        Modem.__init__(self, mm, name,
                       mdn='6175551212',
                       activation_state=Modem.ACTIVATED)

    def ConnectDone(self, old, new, why):
        carrier = self.manager.carrier
        # Implement ConnectDone by manipulating the IP pool restrictor
        if new == flimflam_test.STATE_CONNECTED and carrier.restricted:
            restrictor.enter()
        else:
            restrictor.leave()
        Modem.ConnectDone(self, old, new, why)

    def Connect(self, props, *args, **kwargs):
        print 'ActivatedModem: Connect'
        kwargs['connect_delay_ms'] = (
            self.manager.options.connect_delay_ms)
        Modem.Connect(self, props, *args, **kwargs)

    def ActivateImpl(self, _s, _args, _kwargs):
        raise NotImplementedError('Unimplemented.  Must implement in subclass.')


class BrokenActivatedModem(Modem):
    """ BrokenActivatedModem is a modem that although activated always
        fails to connect to the network.  This simulates errors in which the
        carrier refuses to allow connections.
    """
    def __init__(self, mm, name):
        Modem.__init__(self, mm, name,
                       mdn='6175551212',
                       activation_state=Modem.ACTIVATED)

    # Implement connect by always failing
    def Connect(self, _props, *_args, **_kwargs):
        print 'BrokenActivatedModem: Connect'
        time.sleep(self.manager.options.connect_delay_ms / 1000.0)
        self.state = flimflam_test.STATE_CONNECTING
        glib.timeout_add(500, lambda: self.ConnectDone(
                         self.state,
                         flimflam_test.STATE_REGISTERED,
                         flimflam_test.REASON_USER_REQUESTED))
        raise flimflam_test.ConnectError()

    def ActivateImpl(self, _s, _args, _kwargs):
        raise NotImplementedError('Unimplemented.  Must implement in subclass.')


class Manager(ModemManager):
    def __init__(self, bus, options):
        ModemManager.__init__(self, bus, flimflam_test.OCMM)
        self.modem_number = 1
        self.options = options
        self.carrier = CarrierState(bus,
                                    '/org/chromium/ModemManager/Carrier')

    def NewModem(self, classname):
        # modem registeres itself with mm, so does not disappear
        _ = classname(self, '/TestModem/%d' % self.modem_number)
        self.modem_number += 1

    def MakeFactoryResetModem(self):
        self.NewModem(FactoryResetModem)

    def MakePartiallyActivatedModem(self):
        self.NewModem(PartiallyActivatedModem)

    def MakeActivatedModem(self):
        if not self.options.activatable:
            self.NewModem(PartiallyActivatedModem)
        elif self.options.connectable:
            self.NewModem(ActivatedModem)
        else:
            self.NewModem(BrokenActivatedModem)


def main():
    usage = '''
Run the fake cromo program to simulate different modem and carrier
behaviors.  By default with no arguments the program will simulate a
factory reset modem which needs to be activated and requires the user
to sign up for service.

To test for error cases in which connections to the carrier network
always fail, use the --unconnectable flag.  This is particularly
useful when the initial state is set to activated as in:

   sudo fake-cromo -u -s activated

This can be used to simulate the conditions of crosbug.com/11355

Another simulation that corresponds to many field error conditions is
a device that fails OTASP activation.  This can be simulated by using
the -a flag.  The device should start in either the factory or partial
state.

   sudo fake-cromo -a

To test the re-up process, start the modem in the restricted state with

   sudo fake-cromo -s restricted

One can leave the restricted state by fetching
http://localhost:8080/payment_succeeded.html.  This can be done via
the UI by pressing "Buy Plan", or manually.  On the next reconnect the
user should be out of the restricted IP pool.

If the program is interupted while a partially activated modem is in
the connected state it may leave the iptables set up in a way that
causes all tcp traffic to fail, even when using other network
interfaces.  Fix this by restarting fake cromo and tell it to leave
the restricted IP pool

  sudo fake-cromo -l

'''
    parser = OptionParser(usage=usage)
    parser.add_option('-u', '--unconnectable',
                      action='store_false', dest='connectable',
                      default=True,
                      help='Do not allow modem to connect')
    parser.add_option('-s', '--state', dest='initial_state',
                      type='choice',
                      choices=['factory', 'partial',
                           'activated', 'restricted'],
                      default='factory',
                      help=('Set initial state to factory,'
                        'partial, restricted or activated'))
    parser.add_option('-a', '--activation_fails',
                      action='store_false', dest='activatable',
                      default=True,
                      help='Do not allow modem to activate')
    parser.add_option('-l', '--leave-restricted-pool',
                      action='store_true', dest='leave_restricted_pool',
                      default=False,
                      help='Leave the restricted pool and exit')
    parser.add_option('--connect-delay', type='int',
                      dest='connect_delay_ms',
                      default=flimflam_test.DEFAULT_CONNECT_DELAY_MS,
                      help='time in ms required to connnect')

    (options, args) = parser.parse_args()
    if len(args) != 0:
        parser.error("incorrect number of arguments")

    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
    bus = dbus.SystemBus()
    mm = Manager(bus, options)
    mainloop = gobject.MainLoop()
    print "Running test modemmanager."
    _ = dbus.service.BusName(flimflam_test.CMM, bus)

    if options.leave_restricted_pool:
        restrictor.leave(force=True)
        return 0

    # Choose the type of modem to instantiate...
    if options.initial_state == 'factory':
        mm.MakeFactoryResetModem()
    elif options.initial_state == 'partial':
        mm.MakePartiallyActivatedModem()
    elif options.initial_state == 'activated':
        mm.MakeActivatedModem()
    elif options.initial_state == 'restricted':
        mm.carrier.ConsumePlan()
        mm.MakeActivatedModem()
    else:
        print 'Invalid initial state: %s' % options.initial_state
        return 1

    try:
        mainloop.run()
    finally:
        restrictor.leave(force=True)


if __name__ == '__main__':
    main()
